import argparse
import sys
import gzip
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# ==========================================
# 1. 定义正确的阈值 (百分比 0-100)
# ==========================================
CHH_THRESHOLDS = {
    "CAA": 95.0,
    "CAC": 90.0, "CAT": 90.0,
    "CTA": 90.0, "CTT": 90.0, "CTC": 90.0,
    "CCA": 85.0, "CCT": 85.0, "CCC": 85.0,
    "DEFAULT": 90.0
}

def get_chh_threshold(motif: str) -> float:
    """根据 Motif 返回阈值，默认为 90.0"""
    return CHH_THRESHOLDS.get(motif, CHH_THRESHOLDS["DEFAULT"])

def load_bed_frequencies(bed_file: str) -> dict:
    """
    加载 BED 文件中的甲基化频率。
    返回字典: {(chrom, pos, strand): freq}
    """
    logger.info(f"Loading BED file: {bed_file}")
    bed_data = {}
    
    open_func = gzip.open if bed_file.endswith('.gz') else open
    
    count = 0
    with open_func(bed_file, 'rt') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) < 11:
                continue
            
            chrom = parts[0]
            pos = parts[1]
            strand = parts[5]
            
            # 通常第 10 列是 coverage, 第 11 列是 freq (索引 9 和 10)
            # 或者是第 4, 5 列，取决于格式。
            # 这里假设是标准 Bismark/methylKit 格式：... coverage freq
            # 如果你的 BED 格式不同，请根据实际情况调整索引
            try:
                freq = float(parts[10]) # 假设 freq 在第 11 列
            except ValueError:
                continue
                
            key = (chrom, pos, strand)
            bed_data[key] = freq
            
            count += 1
            if count % 1000000 == 0:
                logger.info(f"Loaded {count} sites from BED...")
                
    logger.info(f"Total sites loaded: {len(bed_data)}")
    return bed_data

def process_dataset(input_file: str, output_file: str, bed_data: dict):
    """
    读取生成的数据集，过滤假阳性，写出新文件
    """
    logger.info(f"Processing dataset: {input_file}")
    
    open_func_in = gzip.open if input_file.endswith('.gz') else open
    open_func_out = gzip.open if output_file.endswith('.gz') else open
    
    total_lines = 0
    kept_lines = 0
    filtered_lines = 0
    
    with open_func_in(input_file, 'rt') as fin, \
         open_func_out(output_file, 'wt') as fout:
        
        for line in fin:
            line_strip = line.strip()
            if not line_strip:
                continue
            
            parts = line_strip.split()
            
            # 按照你提供的格式解析
            # 0: ref_name, 1: pos, 2: strand, 6: k_mer, -2: label
            try:
                chrom = parts[0]
                pos = parts[1]
                strand = parts[2]
                k_mer = parts[6]
                label_str = parts[-2] # 倒数第二个是 label
                
                label = int(label_str)
            except IndexError:
                # 如果行格式不对，跳过或记录
                continue

            total_lines += 1
            
            # ==========================
            # 核心过滤逻辑
            # ==========================
            should_keep = False
            
            if label == 1:
                # 这是一个正样本，必须严格检查
                # 1. 从 k_mer 获取 Motif
                # 假设 k_mer 长度 21，中心点在索引 10
                if len(k_mer) == 21:
                    center_idx = 10
                    motif = k_mer[center_idx : center_idx+3] # 取中心及后两位
                    
                    # 再次确认是 CHH (虽然你的上一版代码可能已经过滤了非CHH，但安全起见)
                    if len(motif) == 3:
                        # 2. 获取阈值
                        threshold = get_chh_threshold(motif)
                        
                        # 3. 查 BED 表获取真实频率
                        bed_key = (chrom, pos, strand)
                        if bed_key in bed_data:
                            real_freq = bed_data[bed_key]
                            
                            # 4. 比较：只有频率 >= 阈值 才保留
                            if real_freq >= threshold:
                                should_keep = True
                        else:
                            # 如果 BED 里找不到这个点（很少见，除非文件不匹配），
                            # 保守起见，可以选择丢弃，或者保留。这里选择丢弃以免引入噪声。
                            should_keep = False
                else:
                    # K-mer 长度不对，数据可能有问题
                    should_keep = False
                    
            elif label == 0:
                # 负样本，通常直接保留
                # 如果你想确保负样本频率真的是 0，也可以查 bed_data
                should_keep = True
            else:
                # 其他 label
                should_keep = False

            # ==========================
            # 写出
            # ==========================
            if should_keep:
                fout.write(line)
                kept_lines += 1
            else:
                filtered_lines += 1
                
            if total_lines % 100000 == 0:
                print(f"Processed: {total_lines}, Kept: {kept_lines}, Filtered: {filtered_lines}", end='\r')

    print(f"\nDone! Total: {total_lines}, Kept: {kept_lines}, Filtered: {filtered_lines}")
    logger.info(f"Result saved to {output_file}")

def main():
    parser = argparse.ArgumentParser(description="Re-filter CHH extraction dataset using correct thresholds.")
    parser.add_argument("-i", "--input", required=True, help="Input dataset file (generated by previous code)")
    parser.add_argument("-o", "--output", required=True, help="Output filtered dataset file")
    parser.add_argument("-b", "--bed", required=True, help="Bisulfite BED file (containing methylation freq)")
    
    args = parser.parse_args()
    
    # 1. 加载 BED
    bed_map = load_bed_frequencies(args.bed)
    
    # 2. 过滤数据
    process_dataset(args.input, args.output, bed_map)

if __name__ == "__main__":
    main()